# author:巇橙
# data：2023.3.29

import pandas as pd
import math


# 定义信息熵函数
def get_Ent(data_arr):
	class_arr = []  # 存储类别
	num_arr = []  # 存储类别的数目
	new_data_arr = set(data_arr)
	for i in new_data_arr:
		sum = 0
		class_arr.append(i)
		for j in range(0, len(data_arr)):
			if i == data_arr[j]:
				sum += 1
		num_arr.append(sum)
	# 对输入的列表的单列元素种类及数量进行统计得到列表class_arr和列表num_arr，两个列表为key-value映射关系
	# class_arr = [a1,a2,a3]
	# num_arr = [2,3,4]
	num_all = 0  # 初始化个数
	Ent = 0  # 初始化信息熵
	for i in num_arr:
		num_all += i
	for i in num_arr:
		Ent -= (i / num_all) * math.log2(i / num_all)
	return Ent


# 定义根据类别划分数据
def class_split(data_arr, pre_arr):
	# 此处的data_arr和pre_arr的关系依然为映射列表
	Ent = get_Ent(pre_arr)
	new_data_arr = set(data_arr)
	for i in new_data_arr:
		sum = 0  # 初始化统计变量
		class_arr = []  # 初始化类别变量
		for j in range(0, len(data_arr)):
			if i == data_arr[j]:
				sum += 1
				class_arr.append(pre_arr[j])
		new_Ent = get_Ent(class_arr)
		Ent -= (sum / len(data_arr)) * new_Ent
	return Ent


# 定义读取数据函数
def get_data(data):
	data = pd.read_excel(data)
	data_A = data.iloc[:7, 1].tolist()
	data_B = data.iloc[:7, 2].tolist()
	data_C = data.iloc[:7, 3].tolist()
	return data_A, data_B, data_C


# 主函数入口
if __name__ == "__main__":
	data_A, data_B, data_C = get_data("data1.xlsx")
	print("A的信息增益为：", round(class_split(data_A, data_C), 3))
	print("B的信息增益为：", round(class_split(data_B, data_C), 3))
